#include "imageclassiffier.h"

#include <QDebug>
#include <QDir>
#include <QFileDialog>

#include "ui_imageclassiffier.h"

//解决中文乱码
#if _MSC_VER >= 1600
#pragma execution_character_set( "utf-8" )
#endif

ImageClassiffier::ImageClassiffier( QWidget* parent ) : QMainWindow( parent ), ui( new Ui::ImageClassiffier )
{
    ui->setupUi( this );
    ui->classifyButton->setGeometry( 50, 50, 150, 30 );
    ui->resultLabel->setGeometry( 50, 100, 300, 30 );
    connect( ui->classifyButton, &QPushButton::clicked, this, &ImageClassiffier::classifyImage );
    // 训练svm模型
    svm = trainSVM( "A:\\00_qtai_learing\\CIFAR-10\\CIFAR-10-images\\train\\" );
    qDebug() << "init ok";
}

ImageClassiffier::~ImageClassiffier()
{
    delete ui;
}

Mat ImageClassiffier::extractHOGFeatures( const Mat& image )
{
    HOGDescriptor      hog( Size( 32, 32 ), Size( 24, 24 ), Size( 8, 8 ), Size( 8, 8 ), 9 );
    std::vector<float> featuresVec;
    hog.compute( image, featuresVec );

    Mat features( 1, static_cast<int>( featuresVec.size() ), CV_32F );
    std::copy( featuresVec.begin(), featuresVec.end(), features.begin<float>() );

    return features;
}

Ptr<SVM> ImageClassiffier::trainSVM( const String& datadir )
{
    std::vector<Mat> images;
    std::vector<int> labels;

    QDir dir_cat( QString::fromStdString( datadir + "cat" ) );
    for ( const auto& entry : dir_cat.entryInfoList( QDir::Files ) )
    {
        qDebug() << "cat: " << entry.filePath();
        Mat src = imread( entry.filePath().toStdString(), IMREAD_GRAYSCALE );
        if ( !src.empty() )
        {
            images.push_back( extractHOGFeatures( src ) );
            labels.push_back( 1 );  // 1: cat
        }
    }
    qDebug() << "cat finish";
    QDir dir_dog( QString::fromStdString( datadir + "dog" ) );
    for ( const auto& entry : dir_dog.entryInfoList( QDir::Files ) )
    {
        qDebug() << "dog: " << entry.filePath();
        Mat src = imread( entry.filePath().toStdString(), IMREAD_GRAYSCALE );
        if ( !src.empty() )
        {
            images.push_back( extractHOGFeatures( src ) );
            labels.push_back( 2 );  // 2: dog
        }
    }
    qDebug() << "dog finish";
    qDebug() << " images.size()=" << images.size();
    // 确保所有特征矩阵具有相同的尺寸
    int rows = images[0].rows;
    int cols = images[0].cols;
    for ( size_t i = 1; i < images.size(); ++i )
    {
        if ( images[i].rows != rows || images[i].cols != cols )
        {
            CV_Error( cv::Error::StsUnmatchedSizes, "All feature matrices must have the same size" );
        }
    }
    // 将std::vector<Mat>合并到一个Mat中...
    Mat mat_tra( images.size(), rows * cols, CV_32F );
    for ( size_t i = 0; i < images.size(); ++i )
    {
        images[i].reshape( 1, 1 ).copyTo( mat_tra.row( i ) );
    }

    Mat trainingLabels = Mat( labels ).reshape( 1, labels.size() );

    Ptr<SVM> svm = SVM::create();
    svm->setType( SVM::C_SVC );
    svm->setKernel( SVM::LINEAR );
    svm->setTermCriteria( TermCriteria( TermCriteria::MAX_ITER, 100, 1e-6 ) );
    svm->train( mat_tra, ROW_SAMPLE, trainingLabels );

    return svm;
}

void ImageClassiffier::classifyImage()
{
    QString filePath = QFileDialog::getOpenFileName( this, "选择图片" );
    if ( !filePath.isEmpty() )
    {
        Mat image = imread( filePath.toStdString(), IMREAD_GRAYSCALE );
        if ( !image.empty() )
        {
            Mat   features = extractHOGFeatures( image );
            float response = svm->predict( features );
            ui->resultLabel->setText( QString( "Classification Result: %1" )
                                          .arg(
                                              response == 1   ? "cat"
                                              : response == 2 ? "dog"
                                                              : "unknown" ) );
        }
    }
}
